from markitdown import MarkItDown
from docx import Document
from docx.shared import Pt, RGBColor, Inches
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.oxml.ns import qn
import os
from fastapi import APIRouter, UploadFile, File
from fastapi.responses import FileResponse
import tempfile
import shutil

def set_font(run, font_size=12, is_bold=False):
    """设置文本的字体格式"""
    # 设置英文字体
    run.font.name = 'Times New Roman'
    # 设置中文字体
    run._element.rPr.rFonts.set(qn('w:eastAsia'), '宋体')
    # 设置字号
    run.font.size = Pt(font_size)
    # 设置加粗
    run.font.bold = is_bold
    # 设置字体颜色为黑色
    run.font.color.rgb = RGBColor(0, 0, 0)

# # 使用示例
# input_file = "inputFiles/output_image_寒武纪2023合并资产负债表-图片式.md"
# convert_md_to_docx(input_file)
def convert_md_to_docx(input_md_path, output_docx_path=None):
    # 初始化 MarkItDown
    md = MarkItDown(enable_plugins=False)
    
    # 转换 Markdown
    result = md.convert(input_md_path)
    
    # 如果没有指定输出路径，则使用输入文件名（改扩展名为.docx）
    if output_docx_path is None:
        output_docx_path = os.path.splitext(input_md_path)[0] + '.docx'
    
    # 创建新的 Word 文档
    doc = Document()
    
    # 获取文本内容并处理
    content = result.text_content
    lines = content.split('\n')
    current_table = None
    
    # 获取输入文件所在目录（用于解析相对路径的图片）
    base_dir = os.path.dirname(os.path.abspath(input_md_path))
    
    for line in lines:
        line = line.strip()
        if not line:
            current_table = None  # 遇到空行重置表格状态
            continue
            
        # 检查是否是图片路径
        if any(ext in line.lower() for ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp']):
            try:
                # 清理路径
                img_path = line.strip('[]"\' ')
                if '(' in line and ')' in line:
                    img_path = line[line.find('(')+1:line.rfind(')')].strip()
                
                # 处理路径
                if not os.path.isabs(img_path):
                    full_path = os.path.join(base_dir, img_path)
                else:
                    full_path = img_path
                
                # 检查文件是否存在
                if os.path.exists(full_path):
                    # 添加图片到文档
                    doc.add_picture(full_path, width=Inches(6))
                    # 获取最后添加的段落（图片段落）并居中
                    last_paragraph = doc.paragraphs[-1]
                    last_paragraph.alignment = WD_ALIGN_PARAGRAPH.CENTER
                    continue
                else:
                    print(f"图片未找到: {full_path}")
            except Exception as e:
                print(f"处理图片时出错: {str(e)}")
            
        if line.startswith('#'):
            # 处理标题
            level = line.count('#')
            text = line.lstrip('#').strip()
            heading = doc.add_heading(level=min(level, 9))
            run = heading.add_run(text)
            
            # 根据标题级别设置格式
            if level == 1:
                set_font(run, font_size=16, is_bold=True)
            elif level == 2:
                set_font(run, font_size=14, is_bold=True)
            else:
                set_font(run, font_size=12, is_bold=True)
                
            heading.alignment = WD_ALIGN_PARAGRAPH.LEFT
            
        elif line.startswith('|'):
            # 处理表格
            cells = [cell.strip() for cell in line.split('|')[1:-1]]  # 去掉首尾的 |
            if cells:  # 确保有单元格内容
                if not current_table:
                    # 创建新表格
                    current_table = doc.add_table(rows=1, cols=len(cells))
                    current_table.style = 'Table Grid'
                    # 填充第一行
                    for i, cell_text in enumerate(cells):
                        cell = current_table.cell(0, i)
                        paragraph = cell.paragraphs[0]
                        run = paragraph.add_run(cell_text)
                        set_font(run, font_size=10.5, is_bold=True)  # 表头加粗
                else:
                    # 添加新行
                    new_row = current_table.add_row()
                    for i, cell_text in enumerate(cells):
                        if i < len(cells):
                            cell = new_row.cells[i]
                            paragraph = cell.paragraphs[0]
                            run = paragraph.add_run(cell_text)
                            set_font(run, font_size=10.5)  # 表格内容
                    
        else:
            # 处理普通段落
            if line.startswith('---'):
                continue  # 跳过分隔线
            paragraph = doc.add_paragraph()
            run = paragraph.add_run(line)
            set_font(run)  # 使用默认格式（12号 Times New Roman + 宋体）
    
    # 保存文档
    doc.save(output_docx_path)
    print(f"\n文档已保存到: {output_docx_path}")

# 创建路由
router = APIRouter(
    prefix="/markitdown",
    tags=["markitdown"]
)

@router.post("/md2docx")
async def convert_markdown_to_docx(file: UploadFile = File(...)):
    """
    将上传的Markdown文件转换为DOCX文件
    """
    # 创建临时文件来保存上传的内容
    with tempfile.NamedTemporaryFile(delete=False, suffix='.md') as temp_md:
        shutil.copyfileobj(file.file, temp_md)
        temp_md_path = temp_md.name

    # 创建临时文件用于保存输出的docx
    with tempfile.NamedTemporaryFile(delete=False, suffix='.docx') as temp_docx:
        temp_docx_path = temp_docx.name

    try:
        # 转换文件
        convert_md_to_docx(temp_md_path, temp_docx_path)
        
        # 返回转换后的文件
        return FileResponse(
            temp_docx_path,
            media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
            filename=file.filename.replace('.md', '.docx')
        )
    finally:
        # 清理临时文件
        os.unlink(temp_md_path)
        os.unlink(temp_docx_path)


